package com.atguigu.day08;

import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.common.typeutils.base.VoidSerializer;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.TwoPhaseCommitSinkFunction;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

import java.io.BufferedWriter;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;

public class Example6 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        env.enableCheckpointing(10 * 1000L);

        env
                .addSource(new CounterSource())
                .addSink(new TransactionalFileSink());

        env.execute();
    }

    public static class TransactionalFileSink extends TwoPhaseCommitSinkFunction<Long, String, Void> {
        private BufferedWriter transactionalWriter;

        public TransactionalFileSink() {
            super(StringSerializer.INSTANCE, VoidSerializer.INSTANCE);
        }

        // 数据流第一条数据输出时，或者每次检查点完成时，都会创建一个新的事务
        @Override
        protected String beginTransaction() throws Exception {
            long timeNow = System.currentTimeMillis();
            int subTaskIdx = getRuntimeContext().getIndexOfThisSubtask();
            String filename = timeNow + "-" + subTaskIdx;
            Path tempFilePath = Paths.get("/home/zuoyuan/filetemp/" + filename);
            Files.createFile(tempFilePath);
            this.transactionalWriter = Files.newBufferedWriter(tempFilePath);
            return filename;
        }

        // 预提交时触发
        @Override
        protected void preCommit(String transaction) throws Exception {
            transactionalWriter.flush();
            transactionalWriter.close();
        }

        // 正式提交时触发
        // 接收到作业管理器发送的检查点完成的通知时，触发调用
        @Override
        protected void commit(String transaction) {
            Path tempFilePath = Paths.get("/home/zuoyuan/filetemp/" + transaction);
            if (Files.exists(tempFilePath)) {
                try {
                    Path commitFilePath = Paths.get("/home/zuoyuan/filetarget/" + transaction);
                    Files.move(tempFilePath, commitFilePath);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }

        // 回滚操作
        @Override
        protected void abort(String transaction) {
            Path tempFilePath = Paths.get("/home/zuoyuan/tempfile/" + transaction);
            if (Files.exists(tempFilePath)) {
                try {
                    Files.delete(tempFilePath);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }

        // 每来一条数据，触发一次调用
        @Override
        protected void invoke(String transaction, Long in, Context context) throws Exception {
            transactionalWriter.write(in + "\n");
        }
    }

    public static class CounterSource
            extends RichParallelSourceFunction<Long>
            implements CheckpointedFunction {
        // 偏移量
        private Long offset = 0L;

        private boolean running = true;

        @Override
        public void run(SourceContext<Long> ctx) throws Exception {
            // 初始化锁
            final Object lock = ctx.getCheckpointLock();

            while (running) {
                synchronized (lock) {
                    // 保证以下两条语句的原子性
                    ctx.collect(offset);
                    offset += 1L;
                }
                Thread.sleep(1000L);
            }
        }

        @Override
        public void cancel() {
            running = false;
        }

        // 算子状态
        // 用来保存刚消费完的偏移量，列表状态变量中只有一个元素
        private ListState<Long> state;

        // 程序启动或者故障重启时，触发执行
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            state = context.getOperatorStateStore().getListState(
                    new ListStateDescriptor<Long>(
                            "state",
                            Types.LONG
                    )
            );

            for (Long l : state.get()) {
                offset = l;
            }
        }

        // 接收到检查点分界线时触发调用
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            state.clear();
            state.add(offset);
        }
    }
}
